// Copyright (c) 2024 Huawei Technologies Co., Ltd.
// openFuyao is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package controller

import (
	"context"
	"fmt"
	"os"
	"path/filepath"

	"k8s.io/apimachinery/pkg/runtime/serializer/json"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	v1 "openfuyao.com/npu-operator/api/v1"
)

const (
	componentPathPrefix   = "/var/lib/npu-operator/components"
	driverComponent       = componentPathPrefix + "/driver"
	ociRuntimeComponent   = componentPathPrefix + "/oci-runtime"
	vcschedulerComponent  = componentPathPrefix + "/volcano/volcano-scheduler"
	vccontrollerComponent = componentPathPrefix + "/volcano/volcano-controller"
	trainerComponent      = componentPathPrefix + "/trainer"
	devicePluginComponent = componentPathPrefix + "/device-plugin"
	nodeDComponent        = componentPathPrefix + "/noded"
	exporterComponent     = componentPathPrefix + "/npu-exporter"
	clusterdComponent     = componentPathPrefix + "/clusterd"
	rsControllerComponent = componentPathPrefix + "/resilience-controller"
	mindIOACPComponent    = componentPathPrefix + "/mindio/mindioacp"
	mindIOTFTComponent    = componentPathPrefix + "/mindio/mindiotft"
)

var (
	componentPaths = []string{
		driverComponent,
		ociRuntimeComponent,
		devicePluginComponent,
		trainerComponent,
		nodeDComponent,
		vccontrollerComponent,
		vcschedulerComponent,
		clusterdComponent,
		rsControllerComponent,
		exporterComponent,
		mindIOTFTComponent,
		mindIOACPComponent,
	}

	isComponentManaged = map[string]func(*NPUClusterPolicyReconciler) bool{
		driverComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.Driver.Managed
		},
		ociRuntimeComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.OCIRuntime.Managed
		},
		trainerComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.Trainer.Managed
		},
		devicePluginComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.DevicePlugin.Managed
		},
		nodeDComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.NodeD.Managed
		},
		vccontrollerComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.VCController.Managed
		},
		vcschedulerComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.VCScheduler.Managed
		},
		clusterdComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.ClusterD.Managed
		},
		rsControllerComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.RSController.Managed
		},
		exporterComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.Exporter.Managed
		},
		mindIOTFTComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.MindIOTFT.Managed
		},
		mindIOACPComponent: func(r *NPUClusterPolicyReconciler) bool {
			return r.instance.Spec.MindIOACP.Managed
		},
	}
)

type component struct {
	name      string
	resources []componentResource
}

func (c *component) loadResource(
	ctx context.Context,
	manifest []byte,
	serializer *json.Serializer,
) error {
	logger := log.FromContext(ctx)
	obj, gvk, err := serializer.Decode(manifest, nil, nil)
	if err != nil {
		logger.V(1).Info("Failed to load resource", "manifest", string(manifest))
		return fmt.Errorf("decode manifest: %w", err)
	}
	res := resourceCreators[*gvk](obj.(client.Object))
	c.resources = append(c.resources, res)
	logger.V(1).Info("Resource loaded", "gvk", res.gvk())
	return nil
}

func (c *component) reconcile(ctx context.Context, r *NPUClusterPolicyReconciler) (v1.ComponentState, error) {
	logger := log.FromContext(ctx)

	var currentState v1.ComponentState
	if cs := r.instance.Status.GetComponentStatus(c.name); cs != nil {
		currentState = cs.State
	} else {
		currentState = v1.ComponentState{Phase: v1.ComponentPending}
	}
	needReconcile := len(c.resources) > 0
	if needReconcile && currentState.Phase != v1.ComponentRunning && currentState.Phase != v1.ComponentUnmanaged {
		deployingState := v1.ComponentState{Phase: v1.ComponentDeploying, Reason: reconcilingReason}
		if err := r.updateStatus(ctx, func(s *v1.NPUClusterPolicyStatus) bool {
			return s.UpdateComponentStatus(c.name, deployingState)
		}); err != nil {
			logger.Error(err, "Failed to update deploying state")
			return deployingState, err
		}
	}
	finalState := v1.ComponentState{Phase: v1.ComponentRunning, Reason: reconciledReason}
	isManaged := isComponentManaged[c.name](r)
	if !isManaged {
		logger.V(1).Info("Component is unmanaged")
		finalState = v1.ComponentState{Phase: v1.ComponentUnmanaged, Reason: componentUnmanagedReason}
	}
	for _, resource := range c.resources {
		state, err := resource.clone().reconcile(ctx, r, isManaged)
		if err != nil {
			if state == nil {
				state = &v1.ComponentState{
					Phase:   v1.ComponentPending,
					Reason:  reconcileFailedReason,
					Message: err.Error(),
				}
			}
			return *state, err
		}
		if state != nil {
			finalState = *state
		}
	}
	logger.Info("Component reconciled",
		"state", finalState.Phase, "reason", finalState.Reason, "message", finalState.Message)

	return finalState, nil
}

func (r *NPUClusterPolicyReconciler) loadComponents(ctx context.Context) error {
	// Only check the number of components, since it's non-trivial to check if all resources of an arbitrary component
	// is loaded successfully
	if len(r.components) != len(componentPaths) {
		for _, path := range componentPaths {
			ctx := log.IntoContext(ctx, log.FromContext(ctx, "component", path))
			if err := r.loadComponent(ctx, path); err != nil {
				return fmt.Errorf("load component: %w", err)
			}
		}
		if err := r.updateStatus(ctx, func(s *v1.NPUClusterPolicyStatus) (changed bool) {
			s.Phase = v1.PolicyNotReady
			for _, component := range r.components {
				if s.GetComponentStatus(component.name) == nil {
					changed = true
					s.ComponentStatuses = append(s.ComponentStatuses, v1.ComponentStatus{
						Name:  component.name,
						State: v1.ComponentState{Phase: v1.ComponentPending, Reason: reconcilingReason},
					})
				}
			}
			return
		}); err != nil {
			log.FromContext(ctx).V(1).Error(nil, err.Error())
		}
	}
	return nil
}

func (r *NPUClusterPolicyReconciler) loadComponent(ctx context.Context, path string) error {
	logger := log.FromContext(ctx)
	ctx = log.IntoContext(ctx, logger)

	manifests := readManifests(ctx, path)
	serializer := json.NewSerializerWithOptions(
		json.DefaultMetaFactory,
		r.Scheme,
		r.Scheme,
		json.SerializerOptions{Yaml: true, Pretty: false, Strict: false},
	)
	// Load component resources
	component := component{name: path}
	for _, manifest := range manifests {
		if err := component.loadResource(ctx, manifest, serializer); err != nil {
			return fmt.Errorf("load resource from %s: %w", path, err)
		}
	}
	r.components = append(r.components, component)

	return nil
}

func (r *NPUClusterPolicyReconciler) reconcileComponents(ctx context.Context) (result ctrl.Result, err error) {
	logger := log.FromContext(ctx)

	var state v1.ComponentState
	allReady := true
	for _, component := range r.components {
		logger := logger.WithValues("component", component.name)
		ctx := log.IntoContext(ctx, logger)
		state, err = component.reconcile(ctx, r)
		if statusErr := r.updateStatus(ctx, func(s *v1.NPUClusterPolicyStatus) bool {
			return s.UpdateComponentStatus(component.name, state)
		}); statusErr != nil {
			log.FromContext(ctx).V(1).Error(nil, statusErr.Error())
		}
		if err != nil {
			err = fmt.Errorf("reconcile component (%s): %w", component.name, err)
			ignoreError(r.setConditionsError(ctx, componentNotReadyReason, err.Error()))
			return
		}
		allReady = allReady && (state.Phase == v1.ComponentRunning || state.Phase == v1.ComponentUnmanaged)
	}

	switch {
	case !allReady:
		msg := "at least one component is not ready"
		logger.V(-1).Info(fmt.Sprintf("NPUClusterPolicy not ready, %s", msg))
		ignoreError(r.setConditionsError(ctx, componentNotReadyReason, msg))
		result.RequeueAfter = notReadyRequeueInterval
	case !r.hasNFDLabels:
		ignoreError(r.setConditionsReady(ctx, noNFDLabelsReason, "no NFD labels found"))
		result.RequeueAfter = noNFDLabelRequeueInterval
	case !r.hasNPUNodes:
		err = r.setConditionsReady(ctx, noNPUNodesReason, "no NPU nodes with supported CRI runtime found")
	default:
		err = r.setConditionsReady(ctx, reconciledReason, "all components have been successfully reconciled")
	}

	return
}

func readManifests(ctx context.Context, dir string) [][]byte {
	logger := log.FromContext(ctx)
	logger.Info("Reading component manifest files")
	manifests := [][]byte{}
	if err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		logger := logger.V(1).WithValues()
		if err != nil {
			logger.Error(err, "Failed to stat file", "path", path)
			return nil
		}
		if !info.IsDir() {
			manifest, err := os.ReadFile(path)
			if err != nil {
				logger.Error(err, "Failed to read file", "path", path)
				return nil
			}
			manifests = append(manifests, manifest)
		}
		return nil
	}); err != nil {
		logger.Error(err, "Failed to list component manifest directory")
	}
	return manifests
}
